package org.yanh.es.utils;

import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.bulk.BackoffPolicy;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.core.TimeValue;

import java.util.concurrent.TimeUnit;

/**
 * 类说明：通过BulkProcess批量将Mysql数据导入ElasticSearch中
 */
@Slf4j
public class BulkProcess {
    public static BulkProcessor getBulkProcessor(RestHighLevelClient client) {

        BulkProcessor bulkProcessor = null;
        try {
            BulkProcessor.Listener listener = new BulkProcessor.Listener() {
                @Override
                public void beforeBulk(long executionId, BulkRequest request) {
                    System.out.println("开始执行批量操作，ID: " + executionId);
                }
                @Override
                public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
                    if (response.hasFailures()) {
                        System.out.println("批量操作完成，ID: " + executionId);
                    }
                }
                @Override
                public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
                    System.out.println("批量操作失败，ID: " + executionId);
                    failure.printStackTrace();
                }
            };
            BulkProcessor.Builder builder = BulkProcessor.builder(((bulkRequest, bulkResponseActionListener) -> {
                bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                bulkRequest.timeout(TimeValue.timeValueSeconds(100));
                client.bulkAsync(bulkRequest, RequestOptions.DEFAULT, bulkResponseActionListener);
            }), listener);
            // 当达到5000个操作时触发批量请求
            builder.setBulkActions(5000);
            // 当达到100MB大小时触发批量请求
            builder.setBulkSize(new ByteSizeValue(300L, ByteSizeUnit.MB));
            // 每5秒触发一次批量请求，无论大小和操作数如何
            builder.setFlushInterval(TimeValue.timeValueSeconds(100L));
            // 设置退避策略，以防服务器过载或拒绝请求
            // 第一个参数：等待多长时间再次尝试  第二个参数是重试次数
            builder.setBackoffPolicy(BackoffPolicy.exponentialBackoff(TimeValue.timeValueMillis(1000), 3));
            // 设置并发请求的数量为10，即同时只有一个批量请求在执行
            builder.setConcurrentRequests(10);
            // 注意点：在这里感觉有点坑，官网样例并没有这一步，而笔者因一时粗心也没注意，在调试时注意看才发现，上面对builder设置的属性没有生效
            bulkProcessor = builder.build();
            return bulkProcessor;
        } catch (Exception e) {
            e.printStackTrace();
            try {
                bulkProcessor.awaitClose(100L, TimeUnit.SECONDS);
                client.close();
            } catch (Exception e1) {
                log.error(e1.getMessage());
            }
        }
        return bulkProcessor;
    }
}
